import numpy as np
import openslide
import easyocr
import os
from PIL import Image, ImageDraw, ImageFont

# --- 1. 初始化 ---

# 指定单个 SVS 文件路径（请修改为你想测试的文件）
single_svs_file = (
    "/mnt/data1/zzy/LNM/CRCpro-wsi/2025_10_ALL/1055683.svs"  # ←← 修改这里！
)

# 初始化 OCR 读取器
print("正在加载 EasyOCR 模型...")
reader = easyocr.Reader(["en"], gpu=True)  # 可根据需要添加语言，如 ['ch_sim', 'en']
print("模型加载完毕。")

# 输出图像文件名
output_image_file = "output_annotated_label.png"

# --- 2. 处理单个 SVS 文件 ---
print(f"\n正在处理: {os.path.basename(single_svs_file)}")

try:
    slide = openslide.OpenSlide(single_svs_file)

    if "label" in slide.associated_images:
        label_image = slide.associated_images["label"]
        print("成功提取 'label' 图像。")

        # 执行 OCR
        label_np = np.array(label_image)  # PIL Image → NumPy array (H, W, C)
        result = reader.readtext(label_np)

        if not result:
            print("OCR 未在该标签中识别到任何文本。")
        else:
            print(f"成功识别到 {len(result)} 个文本。")

            # 转换为可绘制的 PIL 图像（确保是 RGB 模式）
            if label_image.mode != "RGB":
                label_image = label_image.convert("RGB")
            draw = ImageDraw.Draw(label_image)

            # 尝试加载一个字体（如果没有，使用默认字体）
            try:
                font = ImageFont.truetype("DejaVuSans.ttf", size=20)
            except OSError:
                font = ImageFont.load_default()

            # 绘制每个识别结果
            for bbox, text, prob in result:
                # bbox 是 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
                # PIL 的 polygon 需要展平坐标
                flat_bbox = [coord for point in bbox for coord in point]
                draw.polygon(flat_bbox, outline="red", width=3)

                # 在左上角写文本和置信度
                x, y = bbox[0]
                draw.text((x, y - 25), f"{text} ({prob:.2f})", fill="red", font=font)

            # 保存带标注的图像
            label_image.save(output_image_file)
            print(f"带 OCR 标注的图像已保存为: {output_image_file}")

    else:
        print("未找到 'label' 关联图像。")

    slide.close()

except Exception as e:
    print(f"处理文件时发生错误: {e}")
    if "slide" in locals() and slide:
        slide.close()
